import numpy as np
import rclpy
import cv2
from cv_bridge import CvBridge
from rclpy.node import Node
from sensor_msgs.msg import Image
from zorbus_interface.msg import RGBDImage


# ========== 图像发布节点 ==========
class ImageDepthPublisher(Node):
    def __init__(self):
        super().__init__('rgbd_publisher')
        self.publisher_ = self.create_publisher(RGBDImage, 'rgbd_topic', 10)
        timer_period = 1.0
        self.timer = self.create_timer(timer_period, self.timer_callback)
        self.bridge = CvBridge()

    def timer_callback(self):
        rgb_msg = self.bridge.cv2_to_imgmsg(self.get_rgb_image(), encoding='bgr8')
        depth_msg = self.bridge.cv2_to_imgmsg(self.get_depth_image(), encoding='16UC1')
        msg = RGBDImage()
        msg.rgb_image = rgb_msg
        msg.depth_image = depth_msg
        self.publisher_.publish(msg)
        self.get_logger().info('Published RGBD image')

    @staticmethod
    def get_rgb_image():
        """生成一个带有渐变和图形结构的 RGB 图像"""
        height, width = 480, 640
        # 创建水平红色渐变背景
        red_channel = np.tile(np.linspace(50, 200, width, dtype=np.uint8), (height, 1))
        green_channel = np.zeros_like(red_channel)
        blue_channel = 255 - red_channel
        rgb = np.stack((blue_channel, green_channel, red_channel), axis=2)
        # 画一个圆形模拟“物体”
        cv2.circle(rgb, center=(width // 2, height // 2), radius=80, color=(0, 255, 0), thickness=-1)  # 绿色圆
        return rgb

    @staticmethod
    def get_depth_image():
        """生成一个带有中心深度圆的深度图，值为 uint16（单位 mm）"""
        height, width = 480, 640
        depth = np.ones((height, width), dtype=np.uint16) * 8000  # 背景 8 米

        # 创建一个深度“近”的圆（模拟近距离目标）
        center_x, center_y = width // 2, height // 2
        for y in range(height):
            for x in range(width):
                dist = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2)
                if dist < 100:
                    depth[y, x] = 2000 + int(dist * 10)  # 圆心 2 米，边缘递增到 3 米
        return depth


# ========== 主函数 ==========
def main(args=None):
    rclpy.init()
    rgbd_pub = ImageDepthPublisher()
    rclpy.spin(rgbd_pub)
    rgbd_pub.destroy_node()
    rclpy.shutdown()
